import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn import init

class depthwise_separable_conv(nn.Module):
    def __init__(self, ch_in, ch_out):
        super(depthwise_separable_conv, self).__init__()
        self.ch_in = ch_in
        self.ch_out = ch_out
        self.depth_conv = nn.Conv2d(ch_in, ch_in, kernel_size=3, padding=1, groups=ch_in)
        self.point_conv = nn.Conv2d(ch_in, ch_out, kernel_size=1)

    def forward(self, x):
        x = self.depth_conv(x)
        x = self.point_conv(x)
        # print(x.shape)
        return x


class SKConv(nn.Module):
    def __init__(self, channels, branches=2, reduce=2, stride=1, len=64):
        super(SKConv, self).__init__()
        len = max(int(channels // reduce), len)
        self.convs = nn.ModuleList([])
        for i in range(branches):
            self.convs.append(nn.Sequential(
                nn.Conv2d(channels, channels, kernel_size=3, stride=stride, padding=1 + i, dilation=1 + i,
                          bias=False),
                nn.BatchNorm2d(channels),
                nn.ReLU(inplace=True)
            ))
        self.gap = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Sequential(
            nn.Conv2d(channels, len, kernel_size=1, stride=1, bias=False),
            nn.BatchNorm2d(len),
            nn.ReLU(inplace=True)
        )
        self.fcs = nn.ModuleList([])
        for i in range(branches):
            self.fcs.append(
                nn.Conv2d(len, channels, kernel_size=1, stride=1,bias=False)
            )
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = [conv(x) for conv in self.convs]
        x = torch.stack(x, dim=1)
        attention = torch.sum(x, dim=1)
        attention = self.gap(attention)
        attention = self.fc(attention)
        attention = [fc(attention) for fc in self.fcs]
        attention = torch.stack(attention, dim=1)
        attention = self.softmax(attention)
        x = torch.sum(x * attention, dim=1)
        return x



class SKCNNNet(nn.Module):
    # ----------------------------
    # Build the model architecture
    # ----------------------------
    def __init__(self):
        super().__init__()
        conv_layers = []
        # self.bn0 = nn.BatchNorm2d(1)
        # First Convolution Block with Relu and Batch Norm. Use Kaiming Initialization
        self.pre = self._pre(1, 8)
        self.SK = SKConv(8)
        # self.conv1 = nn.Conv2d(8, 32, kernel_size=(3, 3), stride=(1, 1), padding=(2, 2))
        self.conv1 = depthwise_separable_conv(8,16)
        self.relu1 = nn.ReLU()
        self.bn1 = nn.BatchNorm2d(16)
        self.mp1 = nn.MaxPool2d(2)
        # self.dp1 = nn.Dropout(p=0.05)
        # init.kaiming_normal_(self.conv1.weight, a=0.1)
        # self.conv1.bias.data.zero_()
        conv_layers += [self.conv1, self.bn1, self.relu1, self.mp1]

        # Second Convolution Block
        # self.conv2 = nn.Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        self.conv2 = depthwise_separable_conv(16,32)
        self.relu2 = nn.ReLU()
        self.bn2 = nn.BatchNorm2d(32)
        self.mp2 = nn.MaxPool2d(2)
        # self.dp2 = nn.Dropout(p=0.05)
        # init.kaiming_normal_(self.conv2.weight, a=0.1)
        # self.conv2.bias.data.zero_()
        conv_layers += [self.conv2, self.bn2, self.relu2, self.mp2]

        # Third Convolution Block
        # self.conv3 = nn.Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        self.conv3 = depthwise_separable_conv(32,64)
        self.relu3 = nn.ReLU()
        self.bn3 = nn.BatchNorm2d(64)
        # self.dp3 = nn.Dropout(p=0.05)
        # init.kaiming_normal_(self.conv3.weight, a=0.1)
        # self.conv3.bias.data.zero_()
        conv_layers += [self.conv3, self.bn3, self.relu3]

        # Fourth Convolution Block
        # self.conv4 = nn.Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        self.conv4 = depthwise_separable_conv(64,128)
        self.relu4 = nn.ReLU()
        self.bn4 = nn.BatchNorm2d(128)
        # init.kaiming_normal_(self.conv4.weight, a=0.1)
        # self.conv4.bias.data.zero_()
        conv_layers += [self.conv4, self.bn4, self.relu4]

        # Linear Classifier
        self.ap = nn.AdaptiveAvgPool2d(output_size=1)

        # wide features
        # self.wide =nn.Sequential(
        #     nn.Linear(in_features=4, out_features=8),
            # nn.ReLU(),
            # # nn.Linear(in_features=16, out_features=32),
            # # nn.Dropout(0.05)
        # )
        self.lin = nn.Linear(in_features=128, out_features=3)
        # self.lin1 = nn.Linear(in_features=80, out_features=128)
        # Wrap the Convolutional Blocks
        self.conv = nn.Sequential(*conv_layers)
        self.dp = nn.Dropout(p=0.2)

    def _pre(self, input_channel, outchannel):
        pre = nn.Sequential(
            nn.Conv2d(input_channel, outchannel, kernel_size=5, stride=1, padding=3, bias=False),
            nn.BatchNorm2d(outchannel),
            # nn.BatchNorm2d(num_features=8),
            nn.ReLU(inplace=True),
        )
        return pre
    # # ----------------------------
    # Forward pass computations
    # ----------------------------
    def forward(self, x):
        # Run the convolutional blocks
        x = x.unsqueeze(1)
        x = self.pre(x)
        x = self.SK(x)
        x = self.conv(x)

        # Adaptive pool and flatten for input to linear layer
        x = self.ap(x)
        x = self.dp(x)
        x_all = x.view(x.shape[0], -1)

        # add wide features and concat two layers
        # print(x1.size())
        # y = self.wide(y)
        # x_all = torch.cat((x_all, y), dim=1)

        #
        # Linear layer
        x_all = self.lin(x_all)

        # Final output
        return x_all

if __name__ == "__main__":
    input = torch.rand(64, 64, 239)
    model = SKCNNNet()

    y = torch.rand(128, 8)
    outputs = model(input)
    total = sum(p.numel() for p in model.parameters())
    print("Total params: %.2fK" %(total/1e3))